package com.example.alibabanacosdiscoveryclient01.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.*;

import java.util.List;

/**
 * @author:guoq
 * @date:2024/6/27
 * @descripion:
 */
@RestController
@RequestMapping("/ai")
@Tag(name = "向量存储(VectorStore)应用")
public class VectorStoreController {
    //    private  DashScopeCloudStore dashScopeCloudStore;
    @Autowired
    VectorStore vectorStore;

    @Value("classpath:/vectorStore/test.pdf")
    private Resource springAiResource;
//    @Autowired
//    public VectorStoreController(DashScopeCloudStore dashScopeCloudStore) {
//        this.dashScopeCloudStore = dashScopeCloudStore;
//    }

    @PostMapping("/vectorSore/add")
    @Operation(summary = "增加索引")
    public String add() {
        //解析文档
        DocumentReader reader = new PagePdfDocumentReader(springAiResource);
        List<Document> documents = reader.get();
        //分割文档
        List<Document> splitDocuments = new TokenTextSplitter().apply(documents);
        vectorStore.add(splitDocuments);
        return "向量数据新增成功";
    }

    @GetMapping("/vectorSore/search")
    @Operation(summary = "文档简单检索")
    public List<Document> store(@RequestParam(value = "message", defaultValue = "什么是矢量数据库") String message) {
        return vectorStore.similaritySearch(message);
    }
    @GetMapping("/vectorSore/defineSearch")
    @Operation(summary = "文档自定义检索")
    public List<Document> defineSearch(@RequestParam(value = "message", defaultValue = "什么是矢量数据库") String message,
                                       @RequestParam(value = "topK", defaultValue = "2") Integer topK,
                                       @RequestParam(value = "threshold",defaultValue = "0.75"  ) double threshold
                                       ) {
//        SearchRequest searchRequest = SearchRequest.query(message)
//                .withTopK(topK)
//                .withSimilarityThreshold(threshold);
        SearchRequest searchRequest = SearchRequest.builder()
                .query(message)
                .topK(topK)
                .similarityThreshold(threshold)
                .build();
        return vectorStore.similaritySearch(searchRequest);
    }
}
